from numpy.lib.stride_tricks import as_strided
import numpy as np
import pandas as pd
global df
df = pd.DataFrame({
    "close1": [1.3,2.1,3,4,5,6,7,8,9, 10,11,12],
    "close2": [9,8,7,6,5,4,3,2.1,1.5, 0, -1, -2],
    "symbal": ["a","a","a", "a","a","a", "b","b","b", "b","b","b"]
})


def birolling(df, w, **kwargs):
    """
        df: dataframe
        w : window
    """
    v = df.values
    d0, d1 = v.shape
    s0, s1 = v.strides

    a = as_strided(v, (d0 - (w - 1), w, d1), (s0, s0, s1))

    rolled_df = pd.concat({
        row: pd.DataFrame(values, columns=df.columns)
        for row, values in zip(df.index, a)
    })

    return rolled_df.groupby(level=0, **kwargs)


def bi_corr(x):
    r = x["close1"].astype("float32").corr(x["close2"].astype("float32"))
    if np.isnan(r):
        return 0
    else:
        return r

def group_corr(gdf , dff):
    gdf["corr"] = birolling(gdf, 2).apply(bi_corr)
    gdf["corr"] = gdf["corr"].shift(1)
    print(gdf.index)
    print(dff)
    return gdf

# df["close_corr"] = birolling(df, 2).apply(bi_corr)
# df["close_corr"] = df['close_corr'].shift(1)
# print(df)

tt = df.groupby("symbal").apply(group_corr, df)
#tt.reset_index()
print(tt)
print(type(tt))
#df["corr"] = tt
#print(df)


